# -*- coding: UTF-8 -*-
"""
@项目名称：get_chain_class.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-19-10:14
"""
import shutil
from glob import glob
from Bio.PDB import PDBParser, PDBIO, is_aa
import os
from typing import Dict, List, Tuple
from ChemTools.is_protein import is_protein


class PDBChainAnalyzer:
    """
    用于分析PDB文件中各链类型的类
    """

    def __init__(self):
        """初始化PDB解析器"""
        self.parser = PDBParser(QUIET=True)  # 创建解析器并忽略警告信息[2,4](@ref)
        self.water_resnames = ["HOH", "WAT", "H2O", "DOD", "TIP3"]
        self.ion_resnames = ["NA", "K", "CL", "MG", "CA", "ZN", "FE", "CU", "MN", "SO4", "PO4", "NO3", "CO3", "OH"]

    def parse_pdb_file(self, pdb_file_path: str):
        """
        解析PDB文件[1,4](@ref)
        参数:
            pdb_file_path: PDB文件路径
        返回:
            Bio.PDB.Structure对象或None(如果解析失败)
        """
        try:
            if not os.path.exists(pdb_file_path):
                print(f"错误: 文件 '{pdb_file_path}' 不存在")
                return None
            # 解析PDB文件[1,4](@ref)
            structure_id = os.path.basename(pdb_file_path).split('.')[0]
            structure = self.parser.get_structure(structure_id, pdb_file_path)
            return structure
        except Exception as e:
            print(f"解析PDB文件时出错: {str(e)}")
            return None

    def analyze_chain_type(self, chain) -> str:
        """
        分析单条链的类型[2](@ref)
        参数:
            chain: Bio.PDB.Chain对象
        返回:
            链类型的字符串描述
        """
        # 获取链中的所有残基
        residues = list(chain.get_residues())
        # 检查空链
        if len(residues) == 0:
            return "空链"
        # 检查是否为蛋白质链（包含标准氨基酸）[2](@ref)
        if any(is_aa(residue) for residue in residues):
            return "蛋白质"
        # 对于非蛋白质链，进一步分析具体类型
        resnames = [residue.get_resname().strip() for residue in residues]
        # 检查是否全是水分子
        if all(resname in self.water_resnames for resname in resnames):
            return "水分子" if len(resnames) == 1 else "水分子团"
        # 检查是否全是离子
        if all(resname in self.ion_resnames for resname in resnames):
            return "离子" if len(resnames) == 1 else "离子团"
        # 检查是否为混合类型（包含水和离子）
        water_count = sum(1 for resname in resnames if resname in self.water_resnames)
        ion_count = sum(1 for resname in resnames if resname in self.ion_resnames)
        if water_count > 0 and ion_count > 0 and water_count + ion_count == len(resnames):
            return "水离子混合物"
        # 检查是否为常见配体或小分子
        if len(resnames) == 1:
            return "配体"
        else:
            return "未知配体分子"

    def analyze_all_chains(self, structure) -> Dict[str, List[Tuple[str, str]]]:
        """
        分析结构中所有链的类型
        参数:
            structure: Bio.PDB.Structure对象
        返回:
            包含各模型和各链类型信息的字典
        """
        chain_results = {}
        # 遍历所有模型[4,6](@ref)
        for model in structure:
            model_id = model.get_id()
            chain_results[model_id] = []
            # 遍历模型中的所有链
            for chain in model:
                chain_id = chain.get_id()
                chain_type = self.analyze_chain_type(chain)
                chain_results[model_id].append((chain_id, chain_type))
        return chain_results

    def save_chain_as_pdb(self, chain, output_path: str):
        """
        将单条链保存为PDB文件[3](@ref)
        参数:
            chain: Bio.PDB.Chain对象
            output_path: 输出文件路径
        """
        io = PDBIO()
        io.set_structure(chain)
        io.save(output_path)

    def generate_report(self, chain_results: Dict[str, List[Tuple[str, str]]], pdb_file_path: str):
        """
        生成分析报告
        参数:
            chain_results: 链分析结果
            pdb_file_path: 分析的PDB文件路径
        """
        # print("=" * 50)
        # print(f"PDB文件分析报告: {os.path.basename(pdb_file_path)}")
        # print("=" * 50)
        total_chains = 0
        type_count = {}
        for model_id, chains in chain_results.items():
            # print(f"\n模型 {model_id}:")
            for chain_id, chain_type in chains:
                # print(f"  链 {chain_id}: {chain_type}")
                total_chains += 1
                type_count[chain_type] = type_count.get(chain_type, 0) + 1
        # print("\n" + "=" * 50)
        # print("统计摘要:")
        # print(f"总链数: {total_chains}")
        # for chain_type, count in type_count.items():
        #     print(f"{chain_type}: {count}条链")
        # print("=" * 50)


# 示例用法
def get_chain_class(pdb_file_path, output_dir):
    # 创建分析器实例
    analyzer = PDBChainAnalyzer()
    # 解析PDB文件
    structure = analyzer.parse_pdb_file(pdb_file_path)
    if structure is None:
        print("无法解析PDB文件，请检查文件路径和格式")
        return
    os.makedirs(output_dir, exist_ok=True)
    # 分析所有链的类型
    chain_results = analyzer.analyze_all_chains(structure)
    # 生成并显示报告
    analyzer.generate_report(chain_results, pdb_file_path)
    # 可选：保存特定类型的链为单独PDB文件
    for model in structure:
        for chain in model:
            chain_id = chain.get_id()
            chain_type = analyzer.analyze_chain_type(chain)
            # 只保存蛋白质链
            output_path = os.path.join(output_dir, f"chain_{chain_type}_{chain_id}.pdb")
            analyzer.save_chain_as_pdb(chain, output_path)
            print(f"已保存蛋白质链 {chain_id} 到 {output_path}")
    pdb_list = glob(output_dir + '/*.pdb')
    for i, pdb_file_path in enumerate(pdb_list):
        if is_protein(pdb_file_path) and os.path.getsize(pdb_file_path) > 20000:  # 蛋白质文件必须大于10kb
            if '蛋白质' not in pdb_file_path:
                dst_path = pdb_file_path.replace('配体', '蛋白质')
                shutil.move(pdb_file_path, dst_path)
            if ' ' in pdb_file_path:
                shutil.move(pdb_file_path, os.path.join(output_dir, f"chain_蛋白质_{i}.pdb"))
        else:
            if '配体' not in pdb_file_path:
                dst_path = pdb_file_path.replace('蛋白质', '配体')
                shutil.move(pdb_file_path, dst_path)
            if ' ' in pdb_file_path:
                shutil.move(pdb_file_path, os.path.join(output_dir, f"chain_配体_{i}.pdb"))


if __name__ == "__main__":
    pdb_file_path = "/4T/AIDDserver/MoleculeGenerate/GenerateInput/Test/贵州/4NF4/4nf4.pdb"  # 指定PDB文件路径（请替换为你的实际文件路径）
    output_directory = f"{os.path.dirname(pdb_file_path)}/extracted_chains_class"  # 输出目录名称
    get_chain_class(pdb_file_path, output_directory)
